#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
将篡改后的图片插入到PDF中
根据duplication_results.json中的配对信息，将篡改后的图片A'替代图片B的位置
"""

import json
import os
import shutil
from typing import Dict, List, Tuple, Optional
from pathlib import Path

import fitz  # PyMuPDF
from PIL import Image
import io


# ==================== 配置 ====================
REAL_PDF_ROOT = "/home/zhangbo/workspace/aigc/genome/dataset/real/real_pdf"
DUPLICATION_RESULTS_JSON = "/home/haozhuodi/iclr2026/image_duplication/111/duplication_results.json"
OUTPUT_ROOT = "/home/haozhuodi/iclr2026/image_duplication/111/data"
PAGE_PX_PER_PT = 1.0  # 像素到点的转换比例


# ==================== 辅助函数 ====================
def safe_overwrite_pdf(doc: fitz.Document, out_pdf: str) -> None:
    """安全覆盖保存PDF，避免repaired PDF报错"""
    try:
        if hasattr(doc, "can_save_incrementally") and doc.can_save_incrementally():
            doc.save(out_pdf, incremental=True, encryption=fitz.PDF_ENCRYPT_KEEP)
            return
        # 某些旧版没有can_save_incrementally()，先试增量
        doc.save(out_pdf, incremental=True, encryption=fitz.PDF_ENCRYPT_KEEP)
        return
    except Exception:
        tmp = out_pdf + ".tmp"
        doc.save(tmp, garbage=4, deflate=True, encryption=fitz.PDF_ENCRYPT_KEEP)
        try:
            os.replace(tmp, out_pdf)
        except Exception:
            shutil.move(tmp, out_pdf)


def page_pixel_to_pt(bbox_px: List[float]) -> List[float]:
    """将像素坐标转换为点坐标"""
    k = PAGE_PX_PER_PT
    return [bbox_px[0] / k, bbox_px[1] / k, bbox_px[2] / k, bbox_px[3] / k]


def crop_and_resize_to_bytes(ori_img_path: str, out_w: int, out_h: int) -> bytes:
    """将图片裁剪并调整到指定尺寸，返回字节流"""
    with Image.open(ori_img_path) as img:
        # 直接调整到目标尺寸，因为我们选择的都是尺寸相似的图片
        img = img.convert("RGB").resize((max(1, out_w), max(1, out_h)), Image.Resampling.LANCZOS)
        bio = io.BytesIO()
        img.save(bio, format="PNG", optimize=True)
        return bio.getvalue()


def load_duplication_results(json_path: str) -> Dict:
    """加载duplication_results.json"""
    with open(json_path, 'r', encoding='utf-8') as f:
        return json.load(f)


def find_panel_by_id(paper_data: Dict, panel_id: str) -> Optional[Dict]:
    """在论文数据中查找指定的panel"""
    for page in paper_data.get('pages', []):
        for figure in page.get('figures', []):
            for panel in figure.get('panels', []):
                if panel.get('panel_id') == panel_id:
                    return panel
    return None


def process_single_paper(paper_data: Dict, output_dir: str) -> bool:
    """处理单篇论文，将篡改后的图片插入PDF"""
    paper_id = paper_data['paper_id']
    print(f"处理论文 {paper_id}...")
    
    # 1. 复制原始PDF到输出目录
    original_pdf_path = os.path.join(REAL_PDF_ROOT, str(paper_id), f"{paper_id}.pdf")
    output_pdf_path = os.path.join(output_dir, f"{paper_id}.pdf")
    
    if not os.path.exists(original_pdf_path):
        print(f"  警告: 原始PDF不存在: {original_pdf_path}")
        return False
    
    # 复制PDF
    shutil.copy2(original_pdf_path, output_pdf_path)
    print(f"  复制PDF: {original_pdf_path} -> {output_pdf_path}")
    
    # 2. 打开PDF文档
    doc = fitz.open(output_pdf_path)
    
    # 3. 查找所有被篡改的panel
    tampered_panels = []
    for page in paper_data.get('pages', []):
        for figure in page.get('figures', []):
            for panel in figure.get('panels', []):
                if panel.get('panel_level_issues', {}).get('has_issue', False):
                    tampered_panels.append(panel)
    
    if not tampered_panels:
        print(f"  论文 {paper_id} 没有需要处理的篡改panel")
        doc.close()
        return True
    
    print(f"  找到 {len(tampered_panels)} 个被篡改的panel")
    
    # 4. 处理每个被篡改的panel
    processed_count = 0
    for panel in tampered_panels:
        try:
            # 获取篡改信息
            issues = panel.get('panel_level_issues', {}).get('issues', [])
            if not issues:
                continue
            
            issue = issues[0]  # 取第一个issue
            evidence = issue.get('evidence', {})
            original_panel_id = evidence.get('original_panel_id')
            tampered_panel_id = evidence.get('tampered_panel_id')
            
            if not original_panel_id or not tampered_panel_id:
                print(f"    警告: panel {panel['panel_id']} 缺少配对信息")
                continue
            
            # 查找被替换的panel（图片B）的位置信息
            # 注意：tampered_panel_id是图片A'，original_panel_id是图片A
            # 我们需要找到图片B的位置来替换
            # 根据配对规则，如果original_panel_id是A，那么B应该是配对的另一个
            # 我们需要从JSON中找到被篡改的panel，它对应的就是图片B的位置
            target_panel = panel  # 当前被标记为has_issue的panel就是图片B的位置
            
            # 获取篡改后的图片路径
            # 从tampered_panel_id中提取figure_id
            # tampered_panel_id格式: {paper_id}_{page}_{figure}_{panel}
            parts = tampered_panel_id.split('_')
            if len(parts) >= 4:
                figure_id = f"{parts[0]}_{parts[1]}_{parts[2]}"  # {paper_id}_{page}_{figure}
                tampered_image_path = os.path.join(output_dir, "panel", figure_id, 
                                                 f"{tampered_panel_id}_tampered.png")
            else:
                print(f"    警告: tampered_panel_id格式错误: {tampered_panel_id}")
                continue
            
            if not os.path.exists(tampered_image_path):
                print(f"    警告: 篡改后的图片不存在: {tampered_image_path}")
                continue
            
            # 获取bbox信息（使用目标panel的bbox）
            bbox_px = target_panel.get('bbox_page_pixel', [])
            if len(bbox_px) != 4:
                print(f"    警告: panel {target_panel['panel_id']} 的bbox格式错误")
                continue
            
            # 计算覆盖尺寸
            w_overlay = int(round(bbox_px[2] - bbox_px[0]))
            h_overlay = int(round(bbox_px[3] - bbox_px[1]))
            
            if w_overlay <= 0 or h_overlay <= 0:
                print(f"    警告: panel {target_panel['panel_id']} 的bbox尺寸无效")
                continue
            
            # 找到对应的页面（使用目标panel的页面）
            page_index = None
            for i, page in enumerate(paper_data.get('pages', [])):
                for figure in page.get('figures', []):
                    for p in figure.get('panels', []):
                        if p.get('panel_id') == target_panel['panel_id']:
                            page_index = page.get('page_index') - 1  # 转换为0-based索引
                            break
                    if page_index is not None:
                        break
                if page_index is not None:
                    break
            
            if page_index is None or page_index >= len(doc):
                print(f"    警告: 找不到panel {target_panel['panel_id']} 对应的页面")
                continue
            
            # 将像素坐标转换为点坐标
            bbox_pt = page_pixel_to_pt(bbox_px)
            
            # 生成覆盖图片的字节流
            overlay_bytes = crop_and_resize_to_bytes(tampered_image_path, w_overlay, h_overlay)
            
            # 在PDF中插入图片
            page = doc[page_index]
            rect = fitz.Rect(*bbox_pt)
            page.insert_image(rect, stream=overlay_bytes, overlay=True, keep_proportion=False)
            
            print(f"    成功替换: {target_panel['panel_id']} <- {tampered_panel_id} (页面 {page_index + 1})")
            processed_count += 1
            
        except Exception as e:
            print(f"    错误: 处理panel {panel.get('panel_id', 'unknown')} 时出错: {e}")
            continue
    
    # 5. 保存修改后的PDF
    safe_overwrite_pdf(doc, output_pdf_path)
    doc.close()
    
    print(f"  论文 {paper_id} 处理完成，成功替换 {processed_count} 个panel")
    return processed_count > 0


def main():
    """主函数"""
    print("开始将篡改后的图片插入PDF...")
    
    # 加载duplication_results.json
    print("加载duplication_results.json...")
    results_data = load_duplication_results(DUPLICATION_RESULTS_JSON)
    papers = results_data.get('papers', [])
    print(f"加载了 {len(papers)} 篇论文")
    
    # 处理每篇论文
    processed_papers = 0
    total_replacements = 0
    
    for paper_data in papers:
        paper_id = paper_data['paper_id']
        output_dir = os.path.join(OUTPUT_ROOT, f"{paper_id}_Duplication")
        
        if not os.path.exists(output_dir):
            print(f"跳过论文 {paper_id}: 输出目录不存在")
            continue
        
        try:
            if process_single_paper(paper_data, output_dir):
                processed_papers += 1
        except Exception as e:
            print(f"处理论文 {paper_id} 时出错: {e}")
            continue
    
    print(f"\n处理完成！")
    print(f"成功处理 {processed_papers} 篇论文")
    print(f"修改后的PDF保存在: {OUTPUT_ROOT}")


if __name__ == "__main__":
    main()
